#!/usr/bin/env python3
import rospy
import numpy as np
import math
from geometry_msgs.msg import PointStamped, Point
from sensor_msgs.msg import PointCloud2
import sensor_msgs.point_cloud2 as pc2
from nav_msgs.msg import Odometry
import cv2

class FrontierGenerator2D:
    def __init__(self):
        rospy.init_node('frontier_generator_2d', anonymous=True)
        
        # 摄像头参数
        self.camera_fov_horizontal = 114.6  # 水平视场角（度）
        self.camera_range = 5.0  # 摄像头最大探测距离（米）
        self.camera_height = 0.57  # 摄像头高度（米）
        
        # 从ROS参数服务器读取地图参数
        self.map_resolution = rospy.get_param('~map_resolution', 0.5)
        self.map_size = rospy.get_param('~map_size', 1000)
        self.map_center = self.map_size // 2  # 地图中心
        
        # 初始化地图
        self.visited_map = np.zeros((self.map_size, self.map_size), dtype=np.uint8)  # 0:未访问, 1:已访问
        
        # 机器人状态
        self.robot_pos = None
        self.robot_yaw = 0.0
        
        # Frontier相关
        self.frontiers = []
        self.frontier_threshold = 1  # 形成frontier的最小连续未知区域大小
        
        # 可视化控制
        self.visited_map_publish_counter = 0
        self.visited_map_publish_frequency = 5  # 每5次更新发布一次visited_map
        
        # 订阅和发布
        rospy.Subscriber('/magv/odometry/gt', Odometry, self.odom_callback)
        self.frontier_pub = rospy.Publisher('/frontiers_2d', PointCloud2, queue_size=1)
        self.visited_map_pub = rospy.Publisher('/visited_map', PointCloud2, queue_size=1)

        
        # 定时器
        self.update_timer = rospy.Timer(rospy.Duration(0.5), self.update_frontiers)
        
        rospy.loginfo("=" * 50)
        rospy.loginfo("2D Frontier生成器已初始化")
        rospy.loginfo("参数配置:")
        rospy.loginfo(f"  摄像头参数: FOV={self.camera_fov_horizontal}°, 探测距离={self.camera_range}m, 高度={self.camera_height}m")
        rospy.loginfo(f"  地图参数: 分辨率={self.map_resolution}m, 大小={self.map_size}x{self.map_size}")
        rospy.loginfo("=" * 50)

    def odom_callback(self, msg):
        """处理里程计数据"""
        self.robot_pos = np.array([
            msg.pose.pose.position.x,
            msg.pose.pose.position.y,
            msg.pose.pose.position.z
        ])
        
        # 提取机器人朝向（yaw角）
        orientation = msg.pose.pose.orientation
        self.robot_yaw = self.quaternion_to_yaw(orientation)
        
        # 更新访问地图
        self.update_visited_map()



    def quaternion_to_yaw(self, orientation):
        """将四元数转换为yaw角"""
        # 四元数转欧拉角
        x = orientation.x
        y = orientation.y
        z = orientation.z
        w = orientation.w
        
        # 计算yaw角
        yaw = math.atan2(2.0 * (w * z + x * y), 1.0 - 2.0 * (y * y + z * z))
        return yaw

    def world_to_map(self, world_pos):
        """将世界坐标转换为地图坐标"""
        map_x = int((world_pos[0] / self.map_resolution) + self.map_center)
        map_y = int((world_pos[1] / self.map_resolution) + self.map_center)
        return map_x, map_y

    def map_to_world(self, map_pos):
        """将地图坐标转换为世界坐标"""
        world_x = (map_pos[0] -self.map_center) * self.map_resolution
        world_y = (map_pos[1] -self.map_center) * self.map_resolution
        return world_x, world_y

    def update_visited_map(self):
        """更新访问地图，模拟摄像头视野"""
        if self.robot_pos is None:
            return
            
        robot_map_x, robot_map_y = self.world_to_map(self.robot_pos)
        # 检查机器人是否在地图范围内
        if not (0 <= robot_map_x < self.map_size and 0 <= robot_map_y < self.map_size):
            return
            
        # 计算摄像头视野范围
        fov_rad = math.radians(self.camera_fov_horizontal)
        camera_range_pixels = int(self.camera_range / self.map_resolution)
        
        # 在摄像头视野范围内进行射线追踪
        for angle_offset in np.arange(-fov_rad/2, fov_rad/2, 0.1):
            current_angle = self.robot_yaw + angle_offset
            
            for distance in range(1, camera_range_pixels + 1):
                # 计算射线终点
                end_x = robot_map_x + int(distance * math.cos(current_angle))
                end_y = robot_map_y + int(distance * math.sin(current_angle))
                # 检查是否在地图范围内
                if not (0 <= end_x < self.map_size and 0 <= end_y < self.map_size):
                    break
                    
                # 标记为已访问
                self.visited_map[end_y, end_x] = 1

    def generate_frontiers(self):
        """生成frontiers"""
        if self.robot_pos is None:
            return []
            
        frontiers = []
        
        # 使用形态学操作找到frontier区域
        # 1. 找到所有未知区域（未访问的区域）
        unknown_mask = (self.visited_map == 0)
        
        # 2. 找到已访问区域的边界
        visited_mask = self.visited_map == 1
        kernel = np.ones((3, 3), np.uint8)
        visited_dilated = cv2.dilate(visited_mask.astype(np.uint8), kernel, iterations=1)
        
        # 3. Frontier是未知区域与已访问区域边界的交集
        frontier_mask = unknown_mask & (visited_dilated == 1)
        
        # 直接返回所有frontier_mask为True的点的世界坐标
        frontiers = []
        indices = np.argwhere(frontier_mask)
        for y, x in indices:
            world_x, world_y = self.map_to_world((x, y))
            frontiers.append({
                'position': np.array([world_x, world_y, 0.0]),
                'pixel': (x, y)
            })
        return frontiers

    def update_frontiers(self, event):
        """定时更新frontiers"""
        self.frontiers = self.generate_frontiers()
        self.publish_frontiers()
        
        # 控制visited_map的发布频率
        self.visited_map_publish_counter += 1
        if self.visited_map_publish_counter >= self.visited_map_publish_frequency:
            self.publish_visited_map()
            self.visited_map_publish_counter = 0
        

    def publish_frontiers(self):
        """发布frontiers为PointCloud2消息"""
        if not self.frontiers:
            return
            
        # 创建PointCloud2消息
        points = []
        for frontier in self.frontiers:
            points.append([frontier['position'][0], frontier['position'][1], frontier['position'][2]])
        
        # 发布PointCloud2
        header = rospy.Header()
        header.stamp = rospy.Time.now()
        header.frame_id = "map"
        
        cloud = pc2.create_cloud_xyz32(header, points)
        self.frontier_pub.publish(cloud)

    def publish_visited_map(self):
        """发布visited_map为PointCloud2消息"""
        if self.robot_pos is None:
            return
            
        # 创建PointCloud2消息
        points = []
        
        # 使用numpy的where函数找到所有已访问的点，提高效率
        visited_indices = np.where(self.visited_map == 1)
        
        for y, x in zip(visited_indices[0], visited_indices[1]):
            world_x, world_y = self.map_to_world((x, y))
            points.append([world_x, world_y, 0.0])
        
        if points:
            # 发布PointCloud2
            header = rospy.Header()
            header.stamp = rospy.Time.now()
            header.frame_id = "map"
            cloud = pc2.create_cloud_xyz32(header, points)
            self.visited_map_pub.publish(cloud)




if __name__ == '__main__':
    try:
        generator = FrontierGenerator2D()
        
        rospy.spin()
    except rospy.ROSInterruptException:
        pass 